import datetime
import time
from urllib.request import Request
from fastapi import HTTPException
import re

import auth.jwt_tools
import cache.redis_cache
import log
import cache.keys
from filter import FILTERS_IP_WHITELIST, FILTERS_IP_BLACKLIST, \
    FILTERS_RATELIMIT_IP_LIMITS, FILTERS_RATELIMIT_IP_WINDOW, FILTERS_RATELIMIT_SVR_LIMITS, \
    FILTERS_RATELIMIT_SVR_WINDOW, FILTER_CONTENT_TYPE_LIST, FILTER_METHOD_LIST, FILTER_BOT_UA_LIST, \
    FILTER_TRUSTED_REFER_LIST
from utils.alarm_util import send_alarm
from utils.auth_util import getRealClientIp


def init_routes(app):
    @app.middleware("http")
    async def content_type_filter(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if len(FILTER_CONTENT_TYPE_LIST) == 0:
            return await call_next(request)

        content_type = request.headers.get("content-type", "")
        if content_type == "":
            raise HTTPException(status_code=400, detail="Missing Content-Type")
        for contentType in FILTER_CONTENT_TYPE_LIST:
            if contentType.lower() not in content_type.lower():
                raise HTTPException(status_code=400, detail="Content-Type Not Right, " + content_type)

        return await call_next(request)

    @app.middleware("http")
    async def block_bots(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if len(FILTER_BOT_UA_LIST) == 0:
            return await call_next(request)

        ua = request.headers.get("User-Agent", "")
        if ua == "":
            raise HTTPException(status_code=400, detail="Missing User-Agent")
        for key in FILTER_BOT_UA_LIST:
            if key.lower() in ua.lower():
                raise HTTPException(status_code=403, detail="Bot detected")

        return await call_next(request)

    @app.middleware("http")
    async def check_referer(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if len(FILTER_TRUSTED_REFER_LIST) == 0:
            return await call_next(request)

        referer = request.headers.get("referer", "")
        if referer == "":
            raise HTTPException(status_code=400, detail="Missing Referer")

        for r in FILTER_TRUSTED_REFER_LIST:
            if r.lower() in referer.lower():
                return await call_next(request)

        raise HTTPException(status_code=403, detail="Referer Non-Valid")

    @app.middleware("http")
    async def method_filter(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if len(FILTER_METHOD_LIST) == 0:
            return await call_next(request)

        methodList = [w.lower() for w in FILTER_METHOD_LIST]
        if request.method.lower() not in methodList:
            raise HTTPException(status_code=405,
                                detail="Method Not Allowed," + request.method)

        return await call_next(request)

    @app.middleware("http")
    async def whitelist_ips(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if getRealClientIp(request) not in FILTERS_IP_WHITELIST:
            if len(FILTERS_IP_WHITELIST) > 0:
                log.logger.error("client ip is not in the whitelist." + getRealClientIp(request))

                send_alarm("whitelist-ip", "Not In White IP List(" + getRealClientIp(request) + ")",
                           "ai-proxy", datetime.datetime.now(), "⚠️")

                raise HTTPException(status_code=401, detail="Reject Connection(not in the whitelist)")

        response = await call_next(request)

        return response

    @app.middleware("http")
    async def blacklist_ips(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        if getRealClientIp(request) in FILTERS_IP_BLACKLIST:
            log.logger.error("client ip is in the blacklist.(" + getRealClientIp(request) + ")")

            send_alarm("blacklist-ip", "In The Black IP List(" + getRealClientIp(request) + ")",
                       "ai-proxy", datetime.datetime.now(), "⚠️")

            raise HTTPException(status_code=401, detail="Reject Connection(in the blacklist)")

        response = await call_next(request)

        return response

    def get_lua_script():
        return """
        local key = KEYS[1]
        local limit = tonumber(ARGV[1])
        local expire = tonumber(ARGV[2])
        
        local current = redis.call('INCR', key)
        
        if current == 1 then
            redis.call('EXPIRE', key, expire)
        end

        return limit - current
        """

    @app.middleware("http")
    async def rate_limit_svr(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        redis = cache.redis_cache.get_redis_connection()
        script = redis.script_load(get_lua_script())
        remaining = redis.evalsha(script, 1, cache.keys.REDIS_CACHE_RATELIMIT_SVR_PREFIX,
                                  FILTERS_RATELIMIT_SVR_LIMITS,
                                  FILTERS_RATELIMIT_SVR_WINDOW)

        if remaining <= 0:
            log.logger.error("rating-limiter-svr rule is triggered.")

            send_alarm("rating-limiter-svr", "Too Many Request(" + FILTERS_RATELIMIT_SVR_LIMITS +
                       "/" + FILTERS_RATELIMIT_SVR_WINDOW + " second(s))",
                       "ai-proxy", datetime.datetime.now(), "⚠️")

            raise HTTPException(status_code=403, detail="Too Many Request(Rate Limit SVR)")

        log.logger.debug("remaining=" + str(remaining))

        response = await call_next(request)

        return response

    @app.middleware("http")
    async def rate_limit_ip(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        clientIp = getRealClientIp(request)
        if clientIp:
            redis = cache.redis_cache.get_redis_connection()
            script = redis.script_load(get_lua_script())
            remaining = redis.evalsha(script, 1, cache.keys.REDIS_CACHE_RATELIMIT_IP_PREFIX + clientIp,
                                      FILTERS_RATELIMIT_IP_LIMITS,
                                      FILTERS_RATELIMIT_IP_WINDOW)

            log.logger.debug("remaining=" + str(remaining))

            if remaining <= 0:
                log.logger.error("rating-limiter-ip rule is triggered. (" + getRealClientIp(request) + ")")

                send_alarm("rating-limiter-ip", "Too Many Request For " + getRealClientIp(request),
                           "ai-proxy", datetime.datetime.now(), "⚠️")

                raise HTTPException(status_code=403, detail="Too Many Request(Rate Limit IP)")

        else:
            log.logger.error("client ip is null(reject).")

            raise HTTPException(status_code=403, detail="Invalid IP(NULL IP)")

        response = await call_next(request)

        return response

    @app.middleware("http")
    async def detect_sql_injection(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        suspicious_patterns = [
            r"\b(union|select|drop|delete|update|truncate)\b",
            r"--\s*$",
            r";\s*$"
        ]

        if any(re.search(p, request.url.path, re.IGNORECASE) for p in suspicious_patterns):
            log.logger.error("SQL Injection rule is triggered.(" + request.url.path + ")")

            send_alarm("sql-injection", "Invalid Request(" + request.url.path + ")",
                       "ai-proxy", datetime.datetime.now(), "⚠️")

            raise HTTPException(403, "Invalid Request(SQL Inject)")

        return await call_next(request)

    @app.middleware("http")
    async def authenticate(request: Request, call_next):
        if re.search("/api/auth/", str(request.url)) or re.search("/health", str(request.url)):
            if not re.search("/auth/checkToken", str(request.url)):
                return await call_next(request)

        token = request.headers.get("Authorization")
        if not token:
            log.logger.error("Token not found in the header.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="None Auth Format")

        if not re.search("Bearer", token):
            log.logger.error("Token format is not right in the header.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="Invalid Auth Format")
        parts = token.split("Bearer")
        if len(parts) > 1:
            token = parts[1].lstrip()
            if not token:
                log.logger.error("Token format is not right in the header.(" + getRealClientIp(request) + ")")

                raise HTTPException(status_code=401, detail="Invalid Token Format")
        else:
            log.logger.error("Token format is not right in the header.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="No Token")

        appId = request.headers.get("appId")
        if not appId:
            log.logger.error("appId is not in the header.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="Invalid Fields")

        userId = request.headers.get("userId")
        if userId and userId != "":
            appId = appId + ":" + userId

        redis = cache.redis_cache.get_redis_connection()
        cacheToken = redis.get(cache.keys.REDIS_CACHE_AUTH_ACCESS_TOKEN_PREFIX + appId)
        if not cacheToken:
            log.logger.error("Token is not found in the cache.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="Token Not Found")

        if cacheToken != token:
            log.logger.error("Token is not matched in the cache.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="Token Not Matched")

        verifyResult = auth.jwt_tools.verify_jwt_token(cacheToken)
        if verifyResult == 1:
            redis.delete(cache.keys.REDIS_CACHE_AUTH_ACCESS_TOKEN_PREFIX + appId)
            redis.delete(cache.keys.REDIS_CACHE_AUTH_REFRESH_TOKEN_PREFIX + appId)

            raise HTTPException(status_code=401,
                                detail="Auth Failed,Token Expired, Please retry.")
        elif verifyResult >= 2:
            log.logger.error("Token content is not right.(" + getRealClientIp(request) + ")")

            raise HTTPException(status_code=401, detail="Auth Failed,Invalid Token.")

        return await call_next(request)

    @app.middleware("http")
    async def log_requests(request: Request, call_next):
        if re.search("/health", str(request.url)):
            return await call_next(request)

        start_time = int(time.time() * 1000)

        response = await call_next(request)

        end_time = int(time.time() * 1000)
        response_time = end_time - start_time

        log.logger.info(f"{request.method} - {request.url} - {response_time:.2f}ms")

        return response
